#include "decode.hpp"

//Notes:
//vSize = N
//vNodes: 1, 2, ..., N.
//cSize = M
//check Nodes: 1, 2, ..., M

// compute the index of qmatrix
inline long qindex(long& i, long& j){
  return (i * (M+1) + j);
}

// compute the index of rmatrix
inline long rindex(long& i, long& j){
  return (i * (N+1) + j);
}


// initialization of q_ij(0), q_ij(1)s are implied by 1 - q_ij(0).
void init_variable_node(double* qmat, const double* P, const mat_GF2& H){
  for(long v = 1; v <= N; v++){
    const double qij = 1 - P[v];
    for(long c = 1; c <= M; c++){
      if (IsOne(H(c, v))) // then c is a neighbor of v
	  qmat[qindex(v, c)] = qij;
    }
  }
}

// compute response from check node c_j. i.e., r_ji(0)
double comp_check_node_response(const double* qmat, 
			        long& i, long& j, 
				const mat_GF2& H, 
				const double& multresult){
  const double a = 2 * qmat[qindex(i, j)] - 1;
  if (a > 0)
    return (0.5 * (multresult/a) + 0.5);
  double result = 1;
  for(long v = 1; v <= N; v++){
    if ((v != i) && (IsOne(H(j, v)))) {
      result = result * (2 * qmat[qindex(v, j)] - 1);
    }
  }
  return 0.5 * (result + 1);
}

// obtain the neighbors of the variable node v.
vector<long>* get_neighbors_v(const mat_GF2& H, long& v) {
  vector<long>* result = new vector<long>;
  for(long c = 1; c <= M; c++){
    if (IsOne(H(c, v))){
      result -> push_back(c);
    }
  }
  return result;
}

// obtain the neighbors of the check node c.
vector<long>* get_neighbors_c(const mat_GF2& H, long& c) {
  vector<long>* result = new vector<long>;
  for(long v = 1; v <= N; v++){
    if (IsOne(H(c, v))){
      result -> push_back(v);
    }
  }
  return result;
}

// check whether C is a codeword by verifying H * C = 0
inline bool is_code_word(const vec_GF2& C, const mat_GF2& H){
  GF2 x;
  for(long i = 1; i <= M; i++){
    InnerProduct(x, H(i), C);
    if (IsOne(x)) 
      return false;
  }
  return true;
}

inline double multiply_response1(const double* rmat, long& i, long& j, 
			  vector<long>* neighbors, const double& pi){
  double result = 1.0;
  for(long i = 0; i < (long) neighbors -> size(); i++){
    long c = (*neighbors)[i];
    if (c != j)
      result *= rmat[rindex(c, i)];
  }
  return (1-pi) * result;
}

inline double multiply_response2(const double* rmat, long& i, long& j, 
			  vector<long>* neighbors, const double& pi){
  double result = 1.0;
  for(long i = 0; i < (long) neighbors -> size(); i++){
    long c = (*neighbors)[i];
    if (c != j)
      result *= 1 - rmat[rindex(c, i)];
  }
  return pi * result;
}

// compute updates to qmat(i, j) for decoding step 4.
double compute_variable_node_response(const double* rmat, 
				      long& i, long& j, 
				      vector<long>* neighbors, 
				      const double& pi, const double& Q0i, 
				      const double& Q1i){
  if ((pi == 1.0) || (pi == 0.0)) 
    return (1.0 - pi);
  double coef0 = Q0i;
  double coef1 = Q1i;
  double rmatji = rmat[rindex(j, i)];
  if (rmatji == 0.0) {
    coef0 = multiply_response1(rmat, i, j, neighbors, pi);
  }
  else if (rmatji == 1.0) {
    coef1 = multiply_response2(rmat, i, j, neighbors, pi);
  }
  else{
    coef0 = coef0/rmatji;
    coef1 = coef1/(1 - rmatji);
  }
  return (1/(coef0 + coef1)) * coef0;
}

// free the arrays used in the decoding.
inline void free_arrays(double* qmat, double* rmat, double* Q0, double* Q1){
  delete [] qmat;
  delete [] rmat; 
  delete [] Q0;
  delete [] Q1;
}


inline void compute_multiplies(double& multresult, const double* qmat, const mat_GF2& H, long& c){
  multresult = 1;
  double qvc;
  for(long v = 1; v <= N; v++){
    if (IsOne(H(c, v))){
      qvc = qmat[qindex(v, c)];
      if (qvc != 1)
	multresult = multresult * (2 * qvc - 1);
    }
  }
}

inline void compute_rji(double* rmat, const double* qmat, const mat_GF2& H){
  double multresult;
  for(long c = 1; c <= M; c++){
    compute_multiplies(multresult, qmat, H, c);
    for(long v = 1; v <= N; v++){
      if (IsOne(H(c, v))){      
	rmat[rindex(c, v)] = comp_check_node_response(qmat, v, c, H, multresult);
      }
    }
  }
}

// This function implements the LDPC decoder using the sum product
// algorithm in probability domain.
// C: to store the corrected codeword
// H: parity-check matrix .
// iter: to store the number of iterations used in decoding
// P: a list of initial error probabilities.
bool decodeSPA(vec_GF2& C, const mat_GF2& H, int& iter, const double* P){
  iter = 0;
  C.SetLength(N);
  double* qmat = new double[(N+1) * (M+1)];
  double* rmat = new double[(M+1) * (N+1)];
  double* Q0 = new double[N+1];
  double* Q1 = new double[N+1];
  double threshold;

  // 1. initialization of q_ij(0)
  init_variable_node(qmat, P, H);

  while(iter < MAX_ITER){
#ifdef VERBOSE
    cout << "Iteration: " << iter << endl;
#endif
    // 2. compute r_ji(0), where check node c_j and variable node v_i
    // are neighbors.
    compute_rji(rmat, qmat, H);

    // 3. compute the estimation of codewords.
    for(long v = 1; v <= N; v++) {
      vector<long>* neighbors = get_neighbors_v(H, v);
      double prod0 = 1.0 - P[v];
      double prod1 = 1.0 - prod0;
      if (prod0 == 0.0) {
	for(long i = 0; i < (long) neighbors -> size(); i++){
	  long c = (*neighbors)[i];
	  prod1 = prod1 * (1 - rmat[rindex(c, v)]);
	}
      }
      else if (prod1 == 0.0){
	for(long i = 0; i < (long) neighbors -> size(); i++){
	  long c = (*neighbors)[i];
	  prod0 = prod0 * rmat[rindex(c, v)];
	}	
      }
      else {
	for(long i = 0; i < (long) neighbors -> size(); i++){
	  long c = (*neighbors)[i];
	  prod0 = prod0 * rmat[rindex(c, v)];
	  prod1 = prod1 * (1 - rmat[rindex(c, v)]);
	}
      }
      delete neighbors;
      Q0[v] = prod0;
      Q1[v] = prod1;
      threshold = prod1/(prod0 + prod1);
      if (threshold > 0.5)
	C(v) = to_GF2(1);
      else if (threshold < 0.5)
	C(v) = to_GF2(0);
      else{
	cout << threshold << endl;
	cout << "Overflow may happen!" << endl;
	delete neighbors;
	free_arrays(qmat, rmat, Q0, Q1);
	exit(EXIT_FAILURE);
      }
    }

    // if decoding is successful, terminate
    if (is_code_word(C, H)){
#ifdef VERBOSE
	cout << "Success!" << endl;
#endif
      free_arrays(qmat, rmat, Q0, Q1);
      return true;
    }
        
    // 4. variable nodes update the response messages q_ij to check
    // nodes via equation solving.
    for(long v = 1; v <= N; v++){
      vector<long>* neighbors = get_neighbors_v(H, v);
      for(long c = 1; c <= M; c++){
	qmat[qindex(v, c)] = compute_variable_node_response(rmat, v, c, neighbors, P[v], Q0[v], Q1[v]);
      }
      delete neighbors;
    }
    iter++;
  }

#ifdef VERBOSE
    cout << "Warning: Nonconverge!" << endl;
#endif
  free_arrays(qmat, rmat, Q0, Q1);
  return false;
}




